# encoding: utf-8
import json

from flask import Flask
from flask import request

app = Flask(__name__)

import torch
from transformers import BertTokenizer, BertConfig
from torch.utils import data
import numpy as np


class AlbertClassfier(torch.nn.Module):
    def __init__(self, bert_model, bert_config, num_class):
        super(AlbertClassfier, self).__init__()
        self.bert_model = bert_model
        self.dropout = torch.nn.Dropout(0.4)
        self.fc1 = torch.nn.Linear(bert_config.hidden_size, bert_config.hidden_size)
        self.fc2 = torch.nn.Linear(bert_config.hidden_size, num_class)

    def forward(self, token_ids):
        bert_out = self.bert_model(token_ids)[1]  # 句向量 [batch_size,hidden_size]
        bert_out = self.dropout(bert_out)
        bert_out = self.fc1(bert_out)
        bert_out = self.dropout(bert_out)
        bert_out = self.fc2(bert_out)  # [batch_size,num_class]
        return bert_out


model = torch.load('./model1.pt', map_location=lambda storage, loc: storage)
device = torch.device("cuda:0") if torch.cuda.is_available() else 'cpu'
model.to(device)

# text = '''
# 天空，汉语词汇，拼音tiān kōng，谓天际空阔；日月星辰罗列的广大空间。是地球重要的组成部分，在古代中国有天圆地方的说法。人可以在天空观测气象或天文现象，从而得知天气变化、时间的流逝或自己的方位。日出日落可知一日中的时间，晚上月亮的盈亏可以
#
#
#
# '''

pretrained = 'voidful/albert_chinese_small'
tokenizer = BertTokenizer.from_pretrained(pretrained)
config = BertConfig.from_pretrained(pretrained)


def get_train_test_data(string, max_length=100):
    string = ''.join(string.split())
    data = []
    label = []
    ids = tokenizer.encode(string.strip(), max_length=max_length, truncation='longest_first')
    data.append(ids)
    label.append(1)
    return data, label


# dataloader
class DataGen(data.Dataset):
    def __init__(self, data, label):
        self.data = data
        self.label = label

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return np.array(self.data[index]), np.array(self.label[index])


config_ = ['123456']


def predict(text):
    X_train, y_train = get_train_test_data(text)
    train_dataset = DataGen(X_train, y_train)
    train_dataloader = data.DataLoader(train_dataset, batch_size=1)
    with torch.no_grad():
        for step, (token_ids, label) in enumerate(train_dataloader):
            token_ids = token_ids.to(device).long()
            label = label.to(device).long()
            out = model(token_ids)
            print(out)
            zz = list(out)
            print((out.argmax(1) == label).sum().item())
            return int(zz[0][0] > zz[0][1])
            # return int(out.argmax(1))


@app.route('/', methods=['POST'])
def api():
    try:
        text = request.values.get('text', '')
        if len(text) == 0:
            return json.dumps({"code": 200, "msg": 'faile', "result": "文本长度不能为零"})
        result = predict(text)
        return json.dumps({"code": 200, "msg": "success", "result": result})
    except Exception as e:
        return json.dumps({"code": 500, "msg": "error", "result": str(e)})


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=80)
